# AUTHOR: DING
# -*- codeing = utf-8 -*-
# @Time: 2024/2/1 16:35
# @Author: 86139
# @Site: 
# @File: 01-gpu-speedup.py
# @Software: PyCharm

import torch
import time

print(torch.__version__)
print(torch.cuda.is_available())

a = torch.randn(10000, 1000)
b = torch.randn(1000, 2000)
# cpu
t0 = time.time()
c = torch.matmul(a, b)
t1 = time.time()
print(a.device, t1 - t0, c.norm(2))

# gpu
device = torch.device('cuda')
a = a.to(device)
b = b.to(device)

# cuda第一次运行会初始化
t0 = time.time()
c = torch.matmul(a, b)
t2 = time.time()
print(a.device, t2 - t0, c.norm(2))

# 再次计算准确时间
t0 = time.time()
c = torch.matmul(a, b)
t3 = time.time()
print(a.device, t3 - t0, c.norm(2))
